package org.liziyi.huffman.algorithm;

import java.util.*;

/**
 * huffman算法实现
 * 注意: decode，getHuffmanTree,getCharAfterEncode方法调用时要保证encode方法被调用过，且在同一上下文中
 *
 * @author liziyi
 */
public class HuffmanAlgorithm {

    // 记录字符数统计
    private Map<Character, Integer> charStatistic = new HashMap<>();
    // 记录字符映射 k-编码前，v-编码后
    public Map<Character, String> charMap = new HashMap<>();
    // 小根堆，用于构建Huffman树
    private PriorityQueue<Node> priorityQueue = new PriorityQueue<>(Comparator.comparingInt(a -> a.weight));

    /**
     * Huffman树节点
     */
    public static class Node {
        // 节点字符
        public char ch;
        // 权重(反应字符出现的频次)
        public int weight;
        public Node left, right;

        Node() {
        }

        Node(char ch, int weight) {
            this.ch = ch;
            this.weight = weight;
        }

        Node(int weight) {
            this.weight = weight;
        }

    }

    public String encode(String msg) {
        reset();
        countCharacter(msg);
        createHuffmanTree();
        System.out.println(priorityQueue.peek().left+" "+priorityQueue.peek().right);
        encodeWithHuffmanTree(priorityQueue.peek(), new StringBuilder());
        // 按照得到的字符映射进行编码
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < msg.length(); i++) {
            sb.append(charMap.get(msg.charAt(i)));
        }
        return sb.toString();
    }

    public String decode(String msg) {
        return decodeWithHuffmanTree(msg);
    }

    /**
     * 获得Huffman树的树根
     * @return
     */
    public Node getHuffmanTree() {
        return priorityQueue.peek();
    }
    public Map<Character,String> getCharMap(){
        return charMap;
    }
    public Map<Character,Integer> getCharStatistic(){
        return charStatistic;
    }

    /**
     * 获得Huffman树的层数
     * @return
     */
    public int getHuffmanTreeLayer(){
        int layer=0;
        for(String str:charMap.values()){
            layer=Math.max(layer,str.length());
        }
        return layer+1;
    }
    private void countCharacter(String msg) {
        for (int i = 0; i < msg.length(); i++) {
            charStatistic.compute(msg.charAt(i), (k, v) -> v == null ? 1 : v + 1);
        }
    }

    private void createHuffmanTree() {
        charStatistic.forEach((k, v) -> {
            priorityQueue.add(new Node(k, v));
        });
        Node tmp1, tmp2, tmp3;
        while (priorityQueue.size() != 1) {
            tmp1 = priorityQueue.poll();
            tmp2 = priorityQueue.poll();
            tmp3 = new Node(tmp1.weight + tmp2.weight);
            tmp3.left = tmp1;
            tmp3.right = tmp2;
            priorityQueue.add(tmp3);
        }
    }

    /**
     * 遍历Huffman树进行编码（左0，右1）
     *
     * @param node
     * @param sb
     */
    private void encodeWithHuffmanTree(Node node, StringBuilder sb) {
        if (node == null) return;
        if (node.left == null && node.right == null) {
            charMap.put(node.ch, sb.toString());
            return;
        }
        sb.append(0);
        encodeWithHuffmanTree(node.left, sb);
        sb.setCharAt(sb.length()-1,'1');
        encodeWithHuffmanTree(node.right, sb);
        sb.deleteCharAt(sb.length() - 1);
    }

    /**
     * 搜索huffman树进行解码
     *
     * @param msg
     * @return
     */
    private String decodeWithHuffmanTree(String msg) {
        StringBuilder sb = new StringBuilder();
        Node node = priorityQueue.peek();
        for (int i = 0; i < msg.length(); i++) {
            node = (msg.charAt(i) == '0' ? node.left : node.right);
            // 搜索到叶子，解码出一个字符
            if (node.left == null && node.right == null) {
                sb.append(node.ch);
                // 回到根，重新开始搜索
                node = priorityQueue.peek();
            }
        }
        return sb.toString();
    }

    private void reset() {
        charStatistic.clear();
        charMap.clear();
        priorityQueue.clear();
    }
}
